﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using Loki.Resources;

namespace Loki.Utils
{
    /// <summary>
    /// A generic dictionary, which allows both its keys and values 
    /// to be garbage collected if there are no other references
    /// to them than from the dictionary itself.
    /// </summary>
    /// <remarks>
    /// If either the key or value of a particular entry in the dictionary
    /// has been collected, then both the key and value become effectively
    /// unreachable. However, left-over WeakReference objects for the key
    /// and value will physically remain in the dictionary until
    /// RemoveCollectedEntries is called. This will lead to a discrepancy
    /// between the Count property and the number of iterations required
    /// to visit all of the elements of the dictionary using its
    /// enumerator or those of the Keys and Values collections. Similarly,
    /// CopyTo will copy fewer than Count elements in this situation.
    /// </remarks>
    public sealed class WeakDictionary<TKey, TValue> : BaseDictionary<TKey, TValue>
        where TKey : class
        where TValue : class
    {
        private ConcurrentDictionary<WeakReference<TKey>, WeakReference<TValue>> _InternalDictionary;

        private WeakKeyComparer<TKey> _InternalComparer;

        public WeakDictionary()
            : this(0, null)
        {
        }

        public WeakDictionary(int P_Capacity)
            : this(P_Capacity, null)
        {
        }

        public WeakDictionary(IEqualityComparer<TKey> P_Comparer)
            : this(0, P_Comparer)
        {
        }

        public WeakDictionary(int P_Capacity, IEqualityComparer<TKey> P_Comparer)
        {
            this._InternalComparer = new WeakKeyComparer<TKey>(P_Comparer);
            this._InternalDictionary = new ConcurrentDictionary<WeakReference<TKey>, WeakReference<TValue>>(Environment.ProcessorCount * 2, P_Capacity, this._InternalComparer);
        }

        /// <summary>
        /// WARNING: The count returned here may include entries for which
        /// either the key or value objects have already been garbage
        /// collected. Call RemoveCollectedEntries to weed out collected
        ///  entries and update the count accordingly.
        /// </summary>
        public override int Count
        {
            get { return this._InternalDictionary.Count; }
        }

        public override void Add(TKey P_Key, TValue P_Value)
        {
            if (P_Key == null)
            {
                throw LokiServices.Errors.BuildError<ArgumentException>(ErrorMessages.Utils_WeakDictionnay_KeyNullException, Log);
            }

            WeakReference<TKey> L_WeakKey = new WeakKeyReference<TKey>(P_Key, this._InternalComparer);
            WeakReference<TValue> L_WeakValue = WeakReference<TValue>.Create(P_Value);
            this._InternalDictionary.TryAdd(L_WeakKey, L_WeakValue);
        }

        public override bool ContainsKey(TKey P_Key)
        {
            WeakReference<TKey> L_WeakKey = new WeakKeyReference<TKey>(P_Key, this._InternalComparer);
            return this._InternalDictionary.ContainsKey(L_WeakKey);
        }

        public override bool Remove(TKey P_Key)
        {
            WeakReference<TValue> L_WeakValue;
            WeakReference<TKey> L_WeakKey = new WeakKeyReference<TKey>(P_Key, this._InternalComparer);
            return this._InternalDictionary.TryRemove(L_WeakKey, out L_WeakValue);
        }

        public override bool TryGetValue(TKey P_Key, out TValue P_Value)
        {
            WeakReference<TValue> L_WeakValue;
            WeakReference<TKey> L_WeakKey = new WeakKeyReference<TKey>(P_Key, this._InternalComparer);
            if (this._InternalDictionary.TryGetValue(L_WeakKey, out L_WeakValue))
            {
                P_Value = L_WeakValue.Target;
                return L_WeakValue.IsAlive;
            }

            P_Value = null;
            return false;
        }

        protected override void SetValue(TKey P_Key, TValue P_Value)
        {
            WeakReference<TKey> L_WeakKey = new WeakKeyReference<TKey>(P_Key, this._InternalComparer);
            this._InternalDictionary[L_WeakKey] = WeakReference<TValue>.Create(P_Value);
        }

        public override void Clear()
        {
            this._InternalDictionary.Clear();
        }

        public override IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
        {
            foreach (var L_KVP in this._InternalDictionary)
            {
                WeakReference<TKey> L_WeakKey = L_KVP.Key;
                WeakReference<TValue> L_WeakValue = L_KVP.Value;
                TKey L_Key = L_WeakKey.Target;
                TValue L_Value = L_WeakValue.Target;
                if (L_WeakKey.IsAlive && L_WeakValue.IsAlive)
                {
                    yield return new KeyValuePair<TKey, TValue>(L_Key, L_Value);
                }
            }
        }

        /// <summary>
        /// Removes the left-over weak references for entries in the dictionary
        /// whose key or value has already been reclaimed by the garbage
        /// collector. This will reduce the dictionary's Count by the number
        /// of dead key-value pairs that were eliminated.
        /// </summary>
        public void RemoveCollectedEntries()
        {
            List<WeakReference<TKey>> L_ToRemove = null;
            foreach (var L_KVP in this._InternalDictionary)
            {
                WeakReference<TKey> L_WeakKey = L_KVP.Key;
                WeakReference<TValue> L_WeakValue = L_KVP.Value;

                if (!L_WeakKey.IsAlive || !L_WeakValue.IsAlive)
                {
                    if (L_ToRemove == null)
                    {
                        L_ToRemove = new List<WeakReference<TKey>>();
                    }

                    L_ToRemove.Add(L_WeakKey);
                }
            }

            if (L_ToRemove != null)
            {
                WeakReference<TValue> L_WeakValue;
                foreach (var L_Key in L_ToRemove)
                {
                    this._InternalDictionary.TryRemove(L_Key, out L_WeakValue);
                }
            }
        }
    }
}
